"""图像处理工具

提供图像处理、格式转换、尺寸调整等功能。
"""

import os
import io
import base64
from typing import Tuple, Optional, Union, List
from pathlib import Path
from PIL import Image, ImageEnhance, ImageFilter, ImageOps
import numpy as np


class ImageProcessor:
    """图像处理器
    
    提供各种图像处理功能。
    """
    
    # 支持的图像格式
    SUPPORTED_FORMATS = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp'}
    
    @staticmethod
    def load_image(image_path: Union[str, Path]) -> Image.Image:
        """加载图像
        
        Args:
            image_path: 图像路径
            
        Returns:
            PIL图像对象
        """
        try:
            image = Image.open(image_path)
            # 转换为RGB模式（如果不是的话）
            if image.mode != 'RGB':
                image = image.convert('RGB')
            return image
        except Exception as e:
            raise ValueError(f"无法加载图像 {image_path}: {e}")
    
    @staticmethod
    def save_image(
        image: Image.Image,
        output_path: Union[str, Path],
        quality: int = 95,
        optimize: bool = True
    ) -> str:
        """保存图像
        
        Args:
            image: PIL图像对象
            output_path: 输出路径
            quality: JPEG质量（1-100）
            optimize: 是否优化
            
        Returns:
            保存的文件路径
        """
        output_path = Path(output_path)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        
        # 根据文件扩展名确定格式
        format_map = {
            '.jpg': 'JPEG',
            '.jpeg': 'JPEG',
            '.png': 'PNG',
            '.bmp': 'BMP',
            '.tiff': 'TIFF',
            '.webp': 'WEBP'
        }
        
        file_format = format_map.get(output_path.suffix.lower(), 'JPEG')
        
        # 保存参数
        save_kwargs = {'optimize': optimize}
        if file_format == 'JPEG':
            save_kwargs['quality'] = quality
        elif file_format == 'PNG':
            save_kwargs['compress_level'] = 6
        
        try:
            image.save(output_path, format=file_format, **save_kwargs)
            return str(output_path)
        except Exception as e:
            raise ValueError(f"无法保存图像到 {output_path}: {e}")
    
    @staticmethod
    def resize_image(
        image: Image.Image,
        size: Tuple[int, int],
        method: str = 'lanczos',
        maintain_aspect: bool = True
    ) -> Image.Image:
        """调整图像尺寸
        
        Args:
            image: PIL图像对象
            size: 目标尺寸 (width, height)
            method: 重采样方法
            maintain_aspect: 是否保持宽高比
            
        Returns:
            调整后的图像
        """
        # 重采样方法映射
        resample_map = {
            'nearest': Image.NEAREST,
            'bilinear': Image.BILINEAR,
            'bicubic': Image.BICUBIC,
            'lanczos': Image.LANCZOS
        }
        
        resample = resample_map.get(method.lower(), Image.LANCZOS)
        
        if maintain_aspect:
            # 保持宽高比
            image.thumbnail(size, resample)
            return image
        else:
            # 直接调整到指定尺寸
            return image.resize(size, resample)
    
    @staticmethod
    def crop_image(
        image: Image.Image,
        box: Tuple[int, int, int, int]
    ) -> Image.Image:
        """裁剪图像
        
        Args:
            image: PIL图像对象
            box: 裁剪区域 (left, top, right, bottom)
            
        Returns:
            裁剪后的图像
        """
        return image.crop(box)
    
    @staticmethod
    def center_crop(
        image: Image.Image,
        size: Tuple[int, int]
    ) -> Image.Image:
        """中心裁剪
        
        Args:
            image: PIL图像对象
            size: 目标尺寸 (width, height)
            
        Returns:
            裁剪后的图像
        """
        width, height = image.size
        target_width, target_height = size
        
        left = (width - target_width) // 2
        top = (height - target_height) // 2
        right = left + target_width
        bottom = top + target_height
        
        return image.crop((left, top, right, bottom))
    
    @staticmethod
    def pad_image(
        image: Image.Image,
        size: Tuple[int, int],
        color: Tuple[int, int, int] = (255, 255, 255)
    ) -> Image.Image:
        """填充图像到指定尺寸
        
        Args:
            image: PIL图像对象
            size: 目标尺寸 (width, height)
            color: 填充颜色 (R, G, B)
            
        Returns:
            填充后的图像
        """
        return ImageOps.pad(image, size, color=color)
    
    @staticmethod
    def enhance_image(
        image: Image.Image,
        brightness: float = 1.0,
        contrast: float = 1.0,
        saturation: float = 1.0,
        sharpness: float = 1.0
    ) -> Image.Image:
        """图像增强
        
        Args:
            image: PIL图像对象
            brightness: 亮度调整（1.0为原始）
            contrast: 对比度调整（1.0为原始）
            saturation: 饱和度调整（1.0为原始）
            sharpness: 锐度调整（1.0为原始）
            
        Returns:
            增强后的图像
        """
        # 亮度调整
        if brightness != 1.0:
            enhancer = ImageEnhance.Brightness(image)
            image = enhancer.enhance(brightness)
        
        # 对比度调整
        if contrast != 1.0:
            enhancer = ImageEnhance.Contrast(image)
            image = enhancer.enhance(contrast)
        
        # 饱和度调整
        if saturation != 1.0:
            enhancer = ImageEnhance.Color(image)
            image = enhancer.enhance(saturation)
        
        # 锐度调整
        if sharpness != 1.0:
            enhancer = ImageEnhance.Sharpness(image)
            image = enhancer.enhance(sharpness)
        
        return image
    
    @staticmethod
    def apply_filter(
        image: Image.Image,
        filter_type: str,
        **kwargs
    ) -> Image.Image:
        """应用滤镜
        
        Args:
            image: PIL图像对象
            filter_type: 滤镜类型
            **kwargs: 滤镜参数
            
        Returns:
            应用滤镜后的图像
        """
        filter_map = {
            'blur': ImageFilter.BLUR,
            'detail': ImageFilter.DETAIL,
            'edge_enhance': ImageFilter.EDGE_ENHANCE,
            'edge_enhance_more': ImageFilter.EDGE_ENHANCE_MORE,
            'emboss': ImageFilter.EMBOSS,
            'find_edges': ImageFilter.FIND_EDGES,
            'smooth': ImageFilter.SMOOTH,
            'smooth_more': ImageFilter.SMOOTH_MORE,
            'sharpen': ImageFilter.SHARPEN
        }
        
        if filter_type == 'gaussian_blur':
            radius = kwargs.get('radius', 2)
            return image.filter(ImageFilter.GaussianBlur(radius=radius))
        elif filter_type == 'unsharp_mask':
            radius = kwargs.get('radius', 2)
            percent = kwargs.get('percent', 150)
            threshold = kwargs.get('threshold', 3)
            return image.filter(ImageFilter.UnsharpMask(
                radius=radius, percent=percent, threshold=threshold
            ))
        elif filter_type in filter_map:
            return image.filter(filter_map[filter_type])
        else:
            raise ValueError(f"不支持的滤镜类型: {filter_type}")
    
    @staticmethod
    def image_to_base64(image: Image.Image, format: str = 'JPEG') -> str:
        """将图像转换为Base64字符串
        
        Args:
            image: PIL图像对象
            format: 图像格式
            
        Returns:
            Base64编码的字符串
        """
        buffer = io.BytesIO()
        image.save(buffer, format=format)
        img_str = base64.b64encode(buffer.getvalue()).decode()
        return img_str
    
    @staticmethod
    def base64_to_image(base64_str: str) -> Image.Image:
        """将Base64字符串转换为图像
        
        Args:
            base64_str: Base64编码的字符串
            
        Returns:
            PIL图像对象
        """
        img_data = base64.b64decode(base64_str)
        image = Image.open(io.BytesIO(img_data))
        if image.mode != 'RGB':
            image = image.convert('RGB')
        return image
    
    @staticmethod
    def get_image_info(image_path: Union[str, Path]) -> dict:
        """获取图像信息
        
        Args:
            image_path: 图像路径
            
        Returns:
            图像信息字典
        """
        try:
            with Image.open(image_path) as image:
                info = {
                    'path': str(image_path),
                    'format': image.format,
                    'mode': image.mode,
                    'size': image.size,
                    'width': image.width,
                    'height': image.height,
                    'has_transparency': image.mode in ('RGBA', 'LA') or 'transparency' in image.info
                }
                
                # 文件大小
                file_size = os.path.getsize(image_path)
                info['file_size'] = file_size
                info['file_size_mb'] = file_size / (1024 * 1024)
                
                return info
        except Exception as e:
            raise ValueError(f"无法获取图像信息 {image_path}: {e}")
    
    @staticmethod
    def create_grid(
        images: List[Image.Image],
        grid_size: Optional[Tuple[int, int]] = None,
        image_size: Optional[Tuple[int, int]] = None,
        spacing: int = 10,
        background_color: Tuple[int, int, int] = (255, 255, 255)
    ) -> Image.Image:
        """创建图像网格
        
        Args:
            images: 图像列表
            grid_size: 网格尺寸 (cols, rows)，None表示自动计算
            image_size: 单个图像尺寸，None表示使用原始尺寸
            spacing: 图像间距
            background_color: 背景颜色
            
        Returns:
            网格图像
        """
        if not images:
            raise ValueError("图像列表不能为空")
        
        # 自动计算网格尺寸
        if grid_size is None:
            cols = int(np.ceil(np.sqrt(len(images))))
            rows = int(np.ceil(len(images) / cols))
        else:
            cols, rows = grid_size
        
        # 调整图像尺寸
        if image_size is not None:
            images = [ImageProcessor.resize_image(img, image_size, maintain_aspect=False) 
                     for img in images]
        
        # 获取单个图像尺寸
        img_width, img_height = images[0].size
        
        # 计算网格总尺寸
        grid_width = cols * img_width + (cols - 1) * spacing
        grid_height = rows * img_height + (rows - 1) * spacing
        
        # 创建背景图像
        grid_image = Image.new('RGB', (grid_width, grid_height), background_color)
        
        # 粘贴图像
        for i, image in enumerate(images):
            if i >= cols * rows:
                break
            
            row = i // cols
            col = i % cols
            
            x = col * (img_width + spacing)
            y = row * (img_height + spacing)
            
            grid_image.paste(image, (x, y))
        
        return grid_image
    
    @staticmethod
    def validate_image_path(path: Union[str, Path]) -> bool:
        """验证图像路径
        
        Args:
            path: 图像路径
            
        Returns:
            是否为有效的图像路径
        """
        path = Path(path)
        
        # 检查文件是否存在
        if not path.exists():
            return False
        
        # 检查文件扩展名
        if path.suffix.lower() not in ImageProcessor.SUPPORTED_FORMATS:
            return False
        
        # 尝试打开图像
        try:
            with Image.open(path) as image:
                image.verify()
            return True
        except Exception:
            return False


def create_output_filename(
    base_name: str = "generated",
    extension: str = ".jpg",
    timestamp: bool = True,
    counter: Optional[int] = None
) -> str:
    """创建输出文件名
    
    Args:
        base_name: 基础名称
        extension: 文件扩展名
        timestamp: 是否包含时间戳
        counter: 计数器
        
    Returns:
        生成的文件名
    """
    import datetime
    
    filename_parts = [base_name]
    
    if timestamp:
        timestamp_str = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        filename_parts.append(timestamp_str)
    
    if counter is not None:
        filename_parts.append(f"{counter:04d}")
    
    filename = "_".join(filename_parts) + extension
    return filename


def ensure_output_dir(output_dir: Union[str, Path]) -> Path:
    """确保输出目录存在
    
    Args:
        output_dir: 输出目录路径
        
    Returns:
        输出目录路径对象
    """
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    return output_dir